import agg as agg
from pyspark.sql import SparkSession
from pyspark.sql.functions import expr, window, avg
from pyspark.sql.streaming import query
from pyspark.sql.types import *


def parse_data_from_kafka_message(sdf, schema):
    from pyspark.sql.functions import split
    assert sdf.isStreaming == True, "DataFrame doesn't receive streaming data"
    col = split(sdf['value'], ',')
    for idx, field in enumerate(schema):
        sdf = sdf.withColumn(field.name, col.getItem(idx).cast(field.dataType))
    return sdf.select([field.name for field in schema])



if __name__ == '__main__':

    spark = SparkSession.builder \
        .appName("Spark Structured Streaming for taxi ride info") \
        .getOrCreate()

    rides = spark \
        .read \
        .format("parquet").load("spark/stream/taxiData.parquet") \
        .selectExpr("CAST(value AS STRING)")

    fares = spark \
        .read \
        .format("parquet").load("spark/stream/taxiData.parquet")\
        .selectExpr("CAST(value AS STRING)")


    ridesSchema = StructType([
        StructField("VendorID", LongType()), StructField("store_and_fwd_flag", StringType()),
        StructField("tpep_pickup_datetime", TimestampType()), StructField("tpep_dropoff_datetime", TimestampType()),
        StructField("startLon", FloatType()), StructField("startLat", FloatType()),
        StructField("endLon", FloatType()), StructField("endLat", FloatType()),
        StructField("passengerCnt", ShortType()), StructField("taxiId", LongType()),
        StructField("driverId", LongType())])

    faresSchema = StructType([
        StructField("VendorID", LongType()), StructField("taxiId", LongType()),
        StructField("tpep_pickup_datetime", TimestampType()), StructField("tpep_dropoff_datetime", TimestampType()),
        StructField("paymentType", StringType()), StructField("tip", FloatType()),
        StructField("tolls", FloatType()), StructField("totalFare", FloatType())])

    rides = parse_data_from_kafka_message(rides, ridesSchema)
    fares = parse_data_from_kafka_message(fares, faresSchema)

    MIN_LON, MAX_LON, MIN_LAT, MAX_LAT = -73.7, -74.05, 41.0, 40.5
    rides = rides.filter(
        rides["startLon"].between(MIN_LON, MAX_LON) &
        rides["startLat"].between(MIN_LAT, MAX_LAT) &
        rides["endLon"].between(MIN_LON, MAX_LON) &
        rides["endLat"].between(MIN_LAT, MAX_LAT))
    rides = rides.filter(rides["isStart"] == "END")

    faresWithWatermark = fares\
        .selectExpr("rideId AS rideId_fares", "startTime", "totalFare", "tip")\
        .withWatermark("startTime", "30 minutes")

    ridesWithWatermark = rides\
        .selectExpr("rideId", "endTime", "driverId", "taxiId", "startLon", "startLat", "endLon", "endLat")\
        .withWatermark("endTime", "30 minutes")

    joinDF = faresWithWatermark\
        .join(ridesWithWatermark,
              expr("""
              rideId_fares = rideId AND
               endTime > startTime AND
               endTime <= startTime + interval 2 hours
               """))
    tips = joinDF\
        .groupBy(
            window("endTime", "30 minutes", "10 minutes"),
            "area").agg(avg("tip"))

    query.writeStream\
        .outputMode("append")\
        .format("console")\
        .option("truncate", False
                .start()
                .awaitTermination())